package cn.yshujia.aop;

import cn.yshujia.annotation.RateLimiter;
import cn.yshujia.enums.Info;
import cn.yshujia.enums.LimitEnum;
import cn.yshujia.ex.DefineExceptionHandler;
import cn.yshujia.utils.RequestUtils;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletRequest;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import java.util.Collections;
import java.util.List;

/**
 * @author YShuJia
 * @create 2024/4/23
 * @description 限流
 */

@Slf4j
@Aspect
@Component
public class RateLimiterAspect {
	
	
	@Resource (name = "Redis")
	private RedisTemplate<Object, Object> redisTemplate;
	
	@Resource
	private RedisScript<Long> limitScript;
	
	@Before ("@annotation(rateLimiter)")
	public void doBefore(JoinPoint point, RateLimiter rateLimiter) throws DefineExceptionHandler {
		int time = rateLimiter.time();
		int count = rateLimiter.count();
		
		String combineKey = getCombineKey(rateLimiter, point);
		List<Object> keys = Collections.singletonList(combineKey);
		
		Long number = redisTemplate.execute(limitScript, keys, count, time);
		if (number.intValue() > count) {
			throw new DefineExceptionHandler(Info.TRIGGER_LIMIT.getCode(), Info.TRIGGER_LIMIT.getMessage());
		}
		log.info("限制请求'{}',当前请求'{}',缓存key'{}'", count, number.intValue(), combineKey);
	}
	
	public String getCombineKey(RateLimiter rateLimiter, JoinPoint point) {
		String key = rateLimiter.key();
		HttpServletRequest req = RequestUtils.getReq();
		// 获取请求地址 去掉第一个 /
		String url = req.getRequestURI().replaceAll("^/", "");
		// 根据限流类型获取不同的key
		if (LimitEnum.IP.equals(rateLimiter.limitType())) {
			key = key + RequestUtils.getIp() + ":";
		}
		key = key + url.replaceAll("/", "_");
		return key;
	}
}
